package main

import (
	"fmt"
	"regexp"
	"sort"
	"strings"
	"sync"
)

type WordFrequency struct {
	Word      string //单词本身
	Frequency int    //该单词出现的次数
}

func countWords(text string) map[string]int {
	words := splitWords(text)
	wordCount := make(map[string]int)
	for _, word := range words {
		wordCount[word]++
	}
	return wordCount
}

func splitWords(text string) []string {
	re := regexp.MustCompile(`\w+`)                    //可以多个w
	return re.FindAllString(strings.ToLower(text), -1) //使用定义好的正则表达式对象re来搜索和提取文本中的所有匹配（-1）。
}

// 将结果进行排序以及选择前n个频率最高的单词
func topN(wordFrequencies []WordFrequency, n int) []WordFrequency {
	sort.Slice(wordFrequencies, func(i, j int) bool {
		return wordFrequencies[i].Frequency > wordFrequencies[j].Frequency
	})
	if len(wordFrequencies) > n {
		return wordFrequencies[:n]
	}
	return wordFrequencies
}

func countWordsParallel(text string, workers int) map[string]int {
	var wg sync.WaitGroup
	wordCount := make(map[string]int)
	mutex := sync.Mutex{}

	chunkSize := len(text) / workers
	for i := 0; i < workers; i++ {
		start := i * chunkSize
		end := start + chunkSize
		if i == workers-1 {
			end = len(text)
		}
		wg.Add(1)
		go func(start, end int) {
			defer wg.Done()
			chunk := text[start:end]
			chunkCount := countWords(chunk)
			mutex.Lock()
			defer mutex.Unlock()
			for word, freq := range chunkCount {
				wordCount[word] += freq
			}
		}(start, end)
	}
	wg.Wait()
	return wordCount
}

type TrieNode struct {
	children map[rune]*TrieNode
	freq     int
}

func newTrieNode() *TrieNode {
	return &TrieNode{
		children: make(map[rune]*TrieNode),
	}
}

func (t *TrieNode) insert(word string) {
	node := t
	for _, c := range word {
		if _, ok := node.children[c]; !ok {
			node.children[c] = newTrieNode()
		}
		node = node.children[c]
	}
	node.freq++
}

func (t *TrieNode) topN(n int) []WordFrequency {
	var result []WordFrequency
	t.dfs(n, &result)
	return result
}

func (t *TrieNode) dfs(n int, result *[]WordFrequency) {
	if t.freq > 0 {
		*result = append(*result, WordFrequency{
			Word:      string(t.dfsPath()),
			Frequency: t.freq,
		})
	}
	if len(*result) >= n {
		return
	}
	for _, child := range t.children {
		child.dfs(n, result)
	}
}

func (t *TrieNode) dfsPath() []rune {
	var path []rune
	node := t
	for node != nil {
		for c, child := range node.children {
			if child == node {
				path = append([]rune{c}, path...)
				node = child
				break
			}
		}
	}
	return path
}

func main() {
	text := "This is a sample text. This text has some words. Some words appear more than others."
	wordCount := countWordsParallel(text, 4)

	var wordFrequencies []WordFrequency
	for word, freq := range wordCount {
		wordFrequencies = append(wordFrequencies, WordFrequency{
			Word:      word,
			Frequency: freq,
		})
	}

	topWords := topN(wordFrequencies, 4)
	for _, wf := range topWords {
		fmt.Printf("%s: %d\n", wf.Word, wf.Frequency)
	}

	fmt.Println("\nUsing Trie tree:")
	root := newTrieNode()
	for word, freq := range wordCount {
		for i := 0; i < freq; i++ {
			root.insert(word)
		}
	}
	topWordsFromTrie := root.topN(3)
	for _, wf := range topWordsFromTrie {
		fmt.Printf("%s: %d\n", wf.Word, wf.Frequency)
	}
}
